from langchain_ollama import OllamaLLM
from langchain.chains import create_sql_query_chain
from langchain.schema.runnable import RunnablePassthrough
from langchain_community.tools import QuerySQLDatabaseTool
from langchain.prompts import PromptTemplate
from langchain_ollama import OllamaLLM
import os
import logging
import pandas as pd
from decimal import Decimal
import datetime
from src.utils.helpers import HelperFunctions

class SQLGenerator:
    """SQL生成器，基于LangChain和Ollama
    负责初始化LLM模型、创建SQLChain并生成和执行SQL查询
    """
    def __init__(self, db_connector):
        """初始化SQLGenerator
        
        Args:
            db_connector (PostgreSQLConnector): 数据库连接器实例
        """
        self.db_connector = db_connector
        self.llm = self._initialize_llm()
        self.sql_chain = self._create_sql_chain()
        
    def _initialize_llm(self):
        """初始化Ollama LLM模型
        
        Returns:
            OllamaLLM: 初始化后的Ollama模型实例
        """
        try:
            # 尝试初始化Ollama模型
            return OllamaLLM(
                model=os.getenv("OLLAMA_MODEL"),
                base_url=os.getenv("OLLAMA_URL"),
                temperature=0.2,  # 低随机性，确保SQL稳定性
                system_prompt="你是一个postgresql数据库专家，任何情况下解释部分必须用中文，清晰说明查询逻辑。",
                stream=False  # 禁用流式输出，提高响应速度
            )
        except Exception as e:
            logging.error(f"无法初始化Ollama模型: {str(e)}。")
            return None
    
    def _create_sql_chain(self):
        """创建定制化的SQL链
        
        Returns:
            初始化后的SQL查询链实例，支持实际SQL生成和执行或模拟功能
        """
        try:
            # 获取数据库连接
            db = self.db_connector.get_connection()
            
            # 检查数据库连接是否有效
            if db is None:
                logging.error("数据库连接无效")
                return None
            
            # 正常情况：创建SQL查询链
            logging.info("创建SQL查询链")
            
            # 设置生成的SQL查询结果的最大行数，从环境变量SQL_LIMIT获取， 
            sql_limit = int(os.getenv("SQL_LIMIT", "10000000"))  # 从环境变量获取SQL限制参数
            
            # PostgreSQL专用Prompt模板，取配置文件，第二个参数是默认值
            template_env = os.getenv("POSTGRESQL_PROMPT_TEMPLATE")
            #logging.info(f"Prompt模板:{template_env}")
            # 创建提示词模板 - 使用input变量名以匹配create_sql_query_chain函数要求
            prompt = PromptTemplate(
                input_variables=["table_info", "question", "top_k"],
                template=template_env
            )
            
            # 创建SQL查询链，使用自定义提示词模板
            sql_chain = create_sql_query_chain(
                self.llm,
                db,
                k=sql_limit,  # 
                prompt=prompt  # 使用PostgreSQL专用提示词模板
            )
            
            # 创建查询工具
            execute_query = QuerySQLDatabaseTool(db=db)
            
            # 创建完整链：生成SQL -> 执行SQL
            # 恢复原始链构建方式，确保能正确处理input变量
            full_chain = (
                RunnablePassthrough.assign(query=sql_chain)
                | {"sql": lambda x: x["query"], "result": execute_query}
            )
            
            #logging.info(f"SQL查询链创建成功：{full_chain}")
            return full_chain
        except Exception as e:
            logging.error(f"无法创建SQL链: {str(e)}。")
            # 返回模拟链
            return None
            
    def generate_sql(self, question):
        """生成SQL语句
        
        Args:
            question (str): 用户输入的自然语言查询
            
        Returns:
            str: 生成的SQL语句
        """
        if not question or question.strip() == "":
            return "请输入查询问题"
        logging.info(f"待查询内容: {question}")

        try:
            # 仅生成SQL，不执行
            db = self.db_connector.get_connection()
            sql_chain = create_sql_query_chain(self.llm, db)
            generated_sql = sql_chain.invoke({"question": f"/think {question}"})
            logging.info(f"AI生成的SQL: {generated_sql}")
            return generated_sql
        except Exception as e:
            logging.error(f"SQL生成失败: {str(e)}")
            return f"生成SQL时出错: {str(e)}"
    
    def execute_sql(self, sql_query):
        """执行SQL语句
        
        Args:
            sql_query (str): 要执行的SQL语句
            
        Returns:
            cleaned_sql : 清洗后，可执行的sql语句
            pd.DataFrame: 查询结果的DataFrame
        """
        if not sql_query or sql_query.strip() == "":
            return '',pd.DataFrame({"错误": ["请输入SQL语句"]})
        
        try:
            
            logging.info(f"待查询SQL sql_query: {sql_query}")
            cleaned_sql=sql_query[sql_query.index("</think>")+len("</think>"):]
            cleaned_sql=cleaned_sql[cleaned_sql.index("SQLQuery: ")+len("SQLQuery: "):]
            logging.info(f"去除思考过程的SQL: {cleaned_sql}")
            # 移除可能的代码块标记
            cleaned_sql = cleaned_sql.replace("`", "") #.replace("\n", "")
            # 打印清洗后的SQL
            logging.info(f"清洗后的SQL: {cleaned_sql}")

            # 仅执行SQL
            db = self.db_connector.get_connection()
            execute_query = QuerySQLDatabaseTool(db=db)
            result = execute_query.invoke({"query": cleaned_sql})   
            logging.info(f"SQL执行结果: {result}")

            # 使用 eval 函数将字符串去除首尾空白字符后转换为 Python 对象。
            # eval 函数会执行字符串中的 Python 表达式，这里会将字符串转换为对应的列表数据结构。
            data = eval(result.strip())
            if data:  # 先判断数据不为空
                # 转换Decimal为float（不影响其他类型如int、str、datetime）
                converted_data = []
                for row in data:
                    converted_row = []
                    for item in row:
                        # 若为Decimal类型，转换为float；其他类型保持不变
                        if isinstance(item, Decimal):
                            converted_row.append(float(item))
                        else:
                            converted_row.append(item)
                    converted_data.append(tuple(converted_row))  # 转回元组
                    
                #获取列名
                columns = HelperFunctions.get_columns_name(cleaned_sql)
                if len(columns) == 0:
                    # 可以根据列数动态生成列名或匹配预设列名
                    col_count = len(converted_data[0])  # 列数 = 第一条记录的元素数量
                    logging.info(f"检测到数据有 {col_count} 列")
                    # 示例1：自动生成列名（如col1, col2...）
                    columns = [f"col{i+1}" for i in range(col_count)]
                    '''# 示例2：根据列数匹配预设列名（推荐，更具可读性）
                    # 预设列名映射（根据实际业务调整）
                    column_mapping = {
                        2: ["产品名称", "价格"],
                        3: ["产品名称", "价格", "库存"],
                        4: ["产品名称", "价格", "库存", "分类"]
                    }
                    columns = column_mapping.get(col_count, [f"col{i+1}" for i in range(col_count)])
                    '''
                # 创建DataFrame（此时价格字段已为float类型）
                df = pd.DataFrame(converted_data,columns=columns)
                       
                logging.info(f"转换后的DataFrame:\n{df}")
            else:
                logging.error("查询结果为空")
                df = pd.DataFrame({"查询结果为空"})  
            return cleaned_sql,df
        except Exception as e:
            logging.error(f"SQL执行失败: {str(e)}")
            return '',pd.DataFrame({"错误": [f"执行SQL时出错: {str(e)}"]})
    
    